import time
from math import cos, radians
from typing import Literal

import numpy as np

from mozi_ai_x.simulation.side import CSide
from mozi_ai_x.simulation.scenario import CScenario
from mozi_ai_x.simulation.server import MoziServer
from mozi_ai_x.utils.geo import (
    GeoPoint,
    get_degree,
    get_two_point_distance,
    get_point_with_point_bearing_distance,
)
from mozi_ai_x.utils.log import mprint_with_name

from . import etc


mprint = mprint_with_name("Agent Env")


def get_target_point():
    """
    获取目标点
    """
    lat2 = etc.task_end_point["latitude"]
    lon2 = etc.task_end_point["longitude"]
    return lat2, lon2


def get_target_distance(lat, lon):
    """
    获取目标距离
    """
    lat2, lon2 = get_target_point()
    distance = get_two_point_distance(lon, lat, lon2, lat2)
    return distance


def _get_reward_value(task_heading, current_heading, distance):
    """
    获取奖励值
    """
    angel = abs(task_heading - current_heading)
    cos_value = cos(radians(angel))
    if cos_value >= 0:
        reward = (10000 * cos(radians(angel))) / distance
        return reward
    else:
        neg_reward = (distance * cos(radians(angel))) / 10000
        return neg_reward


def get_distance_reward(lat, lon, last_heading, heading_change):
    """
    获取距离奖励值
    """
    lat2, lon2 = get_target_point()
    distance = get_two_point_distance(lon, lat, lon2, lat2)
    task_heading = get_degree(lat, lon, lat2, lon2)
    current_heading = last_heading + heading_change
    return _get_reward_value(task_heading, current_heading, distance)


class BaseEnvironment:
    """

    环境
    """

    def __init__(
        self,
        ip: str,
        port: int,
        platform: Literal["windows", "linux"],
        scenario_path: str,
        simulate_compression: Literal[0, 1, 2, 3, 4, 5, 6, 7, 8],
        duration_interval: int,
        synchronous: bool = True,
        app_mode: int = 3,
        agent_key_event_file: str | None = None,
        platform_mode: Literal["versus", "development", "train", "eval"] = "development",
    ):
        self.server_ip = ip
        self.port = port
        self.platform: Literal["windows", "linux"] = platform
        self.scenario_path = scenario_path
        self.simulate_compression: Literal[0, 1, 2, 3, 4, 5, 6, 7, 8] = simulate_compression
        self.synchronous = synchronous  # True同步, False异步
        self.duration_interval = duration_interval
        self.app_mode = app_mode
        self.agent_key_event_file = agent_key_event_file
        self.platform_mode: Literal["versus", "development", "train", "eval"] = platform_mode
        self.step_count = 0

        self.mozi_server: MoziServer
        self.scenario: CScenario

    async def reset(self):
        """
        重置函数
        主要用途：
        加载想定，
        """
        self.step_count = 0
        await self.mozi_server.load_scenario()
        self.scenario = CScenario(self.mozi_server)
        await self.mozi_server.init_situation(self.scenario, 1)
        await self.mozi_server.set_simulate_compression(self.simulate_compression)  # 花费了巨长时间，得查查
        await self.mozi_server.run_simulate()

    async def step(self):
        if self.step_count == 0:
            await self.mozi_server.run_simulate()
        else:
            await self.mozi_server.run_simulate()
            # await (
            #     self.mozi_server.run_grpc_simulate()
            # )  # TODO: 发送消息: ReturnObj(Hs_GRPCSimRun()) 会报错，不确定是不是版本问题，先不用
        self.step_count += 1
        self.situation = await self.mozi_server.update_situation(self.scenario)
        for side in self.scenario.sides.values():
            side.static_update()
        return self.scenario

    def create_scenario(self):
        """
        建立一个想定对象
        """
        self.scenario = CScenario(self.mozi_server)

    async def connect_mozi_server(self, ip: str | None = None, port: int | None = None):
        """
        功能：连接墨子服务器
        参数：
        返回：
        """
        if ip is None and port is None:
            self.mozi_server = MoziServer(
                self.server_ip,
                self.port,
                self.platform,
                self.scenario_path,
                self.simulate_compression,
                self.synchronous,
                self.platform_mode,
                self.agent_key_event_file,
            )
        elif ip is not None and port is not None:
            self.mozi_server = MoziServer(
                ip,
                port,
                self.platform,
                self.scenario_path,
                self.simulate_compression,
                self.synchronous,
                self.platform_mode,
                self.agent_key_event_file,
            )
        else:
            raise ValueError("请正确配置墨子IP与端口！！！")

        await self.mozi_server.start()
        time.sleep(4.0)

    async def start(self, ip: str | None = None, port: int | None = None):
        """
        开始函数
        主要用途：
            1.连接服务器端
            2.设置运行模式
            3.设置步长参数
        """
        if ip is None and port is None:
            await self.connect_mozi_server()
        elif ip is not None and port is not None:
            await self.connect_mozi_server(ip, port)
        else:
            raise ValueError("请正确配置墨子IP与端口！！！")

        await self.mozi_server.set_run_mode(self.synchronous)
        await self.mozi_server.set_decision_step_length(self.duration_interval)

    def is_done(self):
        """
        判定推演是否结束
        """
        response_dic = self.scenario.get_responses()
        for _, v in response_dic.items():
            if v.type == "EndOfDeduction":
                if self.agent_key_event_file:
                    self.mozi_server.write_key_event_string_to_file("推演结束！")
                return True
        return False


class Environment(BaseEnvironment):
    """
    环境类
    """

    def __init__(
        self,
        ip: str,
        port: int,
        scenario_path: str,
        simulate_compression: Literal[0, 1, 2, 3, 4, 5, 6, 7, 8],
        duration_interval: int,
        server_platform: Literal["windows", "linux"],
    ):
        super(Environment, self).__init__(
            ip,
            port,
            server_platform,
            scenario_path,
            simulate_compression,
            duration_interval,
        )
        self.action_space = 2
        self.action_max = 1
        self.state_space = 3
        self.red_unit_list: list[str] = []
        self.observation: list[float] = []
        self.red_side_name = "美国海军"
        self.blue_side_name = "俄罗斯海军"

        self.redside: CSide
        self.blueside: CSide

        self.agent_key_event_file = None
        self.platform_mode = "development"

    def _init_red_unit_list(self):
        """
        初始化红方单元列表
        """
        ret_lt = []

        aircraft_list_dic = self.redside.aircrafts
        for key in aircraft_list_dic:
            ret_lt.append(key)
        return ret_lt

    def _get_a_side_observation(self, unit_list: list[str]) -> list[float]:
        """
        获取一方的观察
        """

        obs_lt = []
        for key in unit_list:
            aircraft_list_dic = self.redside.aircrafts
            unit = aircraft_list_dic.get(key)
            if unit:
                obs_lt.append(unit.longitude)
                obs_lt.append(unit.latitude)
                obs_lt.append(unit.current_heading)
            else:
                obs_lt.append(0.0)
                obs_lt.append(0.0)
                obs_lt.append(0.0)
        return obs_lt

    def _get_red_observation(self):
        """
        获取红方的观察
        """
        unit_list = self.red_unit_list
        obs_lt = self._get_a_side_observation(unit_list)
        return obs_lt

    def _get_waypoint_heading(self, last_heading, action_value):
        """
        获取航路点朝向
        """

        current_heading = last_heading + action_value
        if current_heading < 0:
            current_heading += 360
        if current_heading > 360:
            current_heading -= 360
        return current_heading

    def _get_new_waypoint(self, heading, lat, lon, distance=20.0):
        """
        获取新的航路点
        """
        point = get_point_with_point_bearing_distance(lat, lon, heading, distance)
        return point

    def _deal_point_data(self, waypoint: GeoPoint):
        """
        处理航路点数据
        """
        lon = str(waypoint["longitude"])
        lat = str(waypoint["latitude"])
        return lon, lat

    def _get_aircraft_waypoint(self, action_value):
        """
        根据人工智能计算所得的动作值，计算飞机的目标航路点
        """
        obs = self.observation
        longitude = obs[0]
        latitude = obs[1]
        heading = obs[2]
        waypoint_heading = self._get_waypoint_heading(heading, action_value[0].item() * 90)
        waypoint = self._get_new_waypoint(waypoint_heading, latitude, longitude)
        distance = get_target_distance(latitude, longitude)
        return waypoint, distance

    def _check_done(self):
        """
        检查是否可以结束
        """
        if not self._check_aircraft_exist():
            return True
        if not self._check_target_exist():
            return True
        return False

    def _check_aircraft_exist(self):
        """
        检查飞机是否存在
        """
        obs = self.observation
        for i in range(len(obs)):
            if obs[i] != 0.0:
                return True
        return False

    def _check_target_exist(self):
        """
        检查目标是否存在
        """
        ret = self.scenario.get_units_by_name(etc.target_name)
        for key in ret:
            ret = self.scenario.unit_is_alive(key)
            if not ret:
                pass
            return ret
        return False

    def _get_target_guid(self) -> str | None:
        """
        获取目标guid
        """
        target_guid = None
        for key in self.blueside.submarines:
            if etc.target_name == self.blueside.submarines[key].name:
                target_guid = key
                return target_guid

        return target_guid

    def _get_contact_target_guid(self):
        if self.redside.contacts:
            for key in self.redside.contacts:
                # dic = self.redside.contacts[key].__dict__
                # actual_guid = self.redside.contacts[key].m_ActualUnit
                for k in self.blueside.submarines:
                    if etc.target_name == self.blueside.submarines[k].name:
                        return key

    def _check_is_contact_target(self):
        if self.redside.contacts:
            for key in self.redside.contacts:
                # dic = self.redside.contacts[key].__dict__
                # actual_guid = self.redside.contacts[key].m_ActualUnit
                for k in self.blueside.submarines:
                    if etc.target_name == self.blueside.submarines[k].name:
                        target_guid = k
                        return target_guid
        return False

    def _check_is_find_target(self):
        """
        检查是否发现目标
        """
        target_guid = self._check_is_contact_target()
        if target_guid:
            return True

        return False

    def _get_distance_reward(self, action_value):
        """
        获取距离奖励
        """
        obs = self.observation
        longitude = obs[0]
        latitude = obs[1]
        heading = obs[2]
        distance = get_target_distance(latitude, longitude)
        action_change_heading = action_value[0].item() * 90
        reward = get_distance_reward(latitude, longitude, heading, action_change_heading)
        return reward, distance

    def _get_reward(self, action_value):
        """
        获取奖励
        """
        reward = 0.0
        if action_value == []:
            return 0.0

        distance_reward, distance = self._get_distance_reward(action_value)
        reward += distance_reward

        if distance < etc.target_radius:
            if self._check_is_drop_sonobuoy(action_value):
                reward += 10.0
        else:
            if not self._check_aircraft_exist():
                reward += -100.0

        if self._check_is_find_target():
            reward += 10.0

        if not self._check_target_exist():
            # pylog.info("destroy target get 150.0 reward")
            reward += 150.0

        return reward

    def _check_is_drop_sonobuoy(self, action_value):
        if action_value[1].item() > 0:
            return True
        return False

    async def _execute_action(self, action_value):
        """
        根据当前的动作值，执行动作
        """

        # 根据动作值计算飞机的下一路径点坐标，以及飞机到坦克的距离
        waypoint, distance = self._get_aircraft_waypoint(action_value)
        airs = self.redside.aircrafts
        for guid in airs:
            aircraft = airs[guid]
            mprint(aircraft.name)

            if self._check_is_find_target():
                contact_target_guid = self._get_contact_target_guid()
                if contact_target_guid:
                    await aircraft.auto_attack(contact_target_guid)
                else:
                    mprint.warning("没有找到目标")
            else:
                if distance < etc.target_radius:
                    if self._check_is_drop_sonobuoy(action_value):
                        mprint.info("drop_sonobuoy")
                        await aircraft.drop_sonobuoy("shallow", "active")

                logitude, latitude = self._deal_point_data(waypoint)
                await aircraft.set_waypoint(float(logitude), float(latitude))

    async def _update(self):
        """
        更新
        """
        # await self.mozi_server.update_situation(self.scenario)  # BaseEnvironment 中已经调用
        self.redside.static_update()
        self.blueside.static_update()

    def _get_observation(self):
        """
        获取观察
        """
        red_obs_lt = self._get_red_observation()
        self.observation = red_obs_lt

        if self._check_aircraft_exist():
            if etc.TRANS_DATA:
                lt = []
                lt.append(red_obs_lt[0] - etc.task_end_point["longitude"])
                lt.append(red_obs_lt[1] - etc.task_end_point["latitude"])
                lt.append(red_obs_lt[2] / 360)
                return lt

        return red_obs_lt

    def _construct_side_entity(self):
        """
        构造各方实体
        """
        if self.scenario is None:
            raise ValueError("scenario not set")
        _redside = self.scenario.get_side_by_name(self.red_side_name)
        assert _redside is not None
        self.redside = _redside
        self.redside.static_construct()
        _blueside = self.scenario.get_side_by_name(self.blue_side_name)
        assert _blueside is not None
        self.blueside = _blueside
        self.blueside.static_construct()

    def _init_unit_list(self):
        """
        初始化单元列表
        """
        self.red_unit_list = self._init_red_unit_list()

    def _get_timesteps(self, action):
        """
        获取单步数据
        """
        obs = self._get_observation()
        reward = self._get_reward(action)
        done = self._check_done()
        info = ""
        return np.array(obs), reward, done, info

    async def step(self):
        """
        输入动作
        执行动作
        更新态势
        获取观察
        获取reward
        检查是否结束
        """
        scenario = await super(Environment, self).step()
        # self._execute_action(action)
        await self._update()
        # await self.mozi_server.run_simulate()
        return scenario

    async def reset(self):
        """
        重置
        """
        await super(Environment, self).reset()
        await self.scenario.set_cur_side_and_dir_view("美国海军", False)  # 设置当前推演方及是否显示导演视图
        self._construct_side_entity()

        self._init_unit_list()
